import torch as pt
import numpy as np
import os
import sys
import pickle
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.preprocessing import MinMaxScaler

# 指定随机种子
np.random.seed(777)
pt.manual_seed(777)

# 2.	使用pytorch完成股票预测（每题5分，共40分）
# ①	数据处理
# 1)	加载data-02-stock_daily数据集（5分）
df = pd.read_csv('./data/data-02-stock_daily.csv', header=1)  # (732, 5)

# 2)	合理定义参数（5分）
ALPHA = 0.001
N_ITERS = 80  # 为了快速演示，这里设定偏小，正式代码请酌情加大
N_STEPS = 7
N_RNN_HIDDEN = 128
spr = 1  # subplot rows
spc = 2  # subplot columns
spn = 0  # subplot index
plt.figure(figsize=(12, 6))   # plot size

# 3)	对数据进行合理预处理，使用前七天的数据预测第八天的(提示：将数据倒序) （5分）
# load
df = df[::-1].to_numpy()
df = MinMaxScaler().fit_transform(df)
# Tidy
M, N_INPUT = df.shape
x = []
y = []
for i in range(M - N_STEPS - 1):
    x.append(df[i:i + N_STEPS])
    y.append(df[i + N_STEPS, -1])
x = np.array(x)
y = np.array(y).reshape(-1, 1)
print('x', x.shape)  # (724, 7, 5)
print('y', y.shape)  # (724,)
# split
x_train, x_test, y_train, y_test = train_test_split(x, y,
                                                    train_size=0.7,
                                                    shuffle=False)
x_train = pt.from_numpy(x_train).float()
y_train = pt.from_numpy(y_train).float()
x_test = pt.from_numpy(x_test).float()
y_test = pt.from_numpy(y_test).float()


# ②	模型创建
# 1)	定义循环类，在类内加入两层长短期记忆（5分）
class MyRnn(pt.nn.Module):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.rnn = pt.nn.LSTM(N_INPUT, N_RNN_HIDDEN,
                              num_layers=2,
                              batch_first=True)
        self.fc1 = pt.nn.Linear(N_RNN_HIDDEN, 64)
        self.relu1 = pt.nn.ReLU()
        self.fc2 = pt.nn.Linear(64, 1)
        self.relu2 = pt.nn.ReLU()

    # 2)	进行正向传播（5分）
    def forward(self, x):
        x, _ = self.rnn(x)
        x = x[:, -1, :]
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        x = self.relu2(x)
        return x


# ③	模型预测
# 1)	进行循环，合理选择损失函数、优化器及学习率（5分）
model = MyRnn()
optim = pt.optim.Adam(params=model.parameters(), lr=ALPHA)
criterion = pt.nn.MSELoss()


def r2(h, y):
    v1 = ((h - y) ** 2).mean()
    v2 = ((y - y.std()) ** 2).mean()
    return 1 - v1 / v2


GROUP = int(np.ceil(N_ITERS / 10))
cost_history = np.zeros(N_ITERS)
for i in range(N_ITERS):
    model.train(True)
    optim.zero_grad()
    h = model(x_train)
    cost = criterion(h, y_train)
    cost.backward()
    optim.step()
    model.train(False)
    cost = cost.detach().numpy()
    r2v = r2(h, y_train).detach().numpy()
    cost_history[i] = cost
    if i % GROUP == 0:
        print(f'#{i + 1}: cost = {cost}, r2 = {r2v}')
if i % GROUP != 0:
    print(f'#{i + 1}: cost = {cost}, r2 = {r2v}')

# 2)	绘制损失值变化曲线（5分）
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(cost_history)
plt.title('Cost values')
plt.xlabel('Iterations')

# 3)	绘制测试集真实值与预测值对比图（5分）
model.eval()
h = model(x_test).detach().numpy().ravel()
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(h, label='Hypothesis', alpha=0.5)
plt.plot(y_test.ravel(), label='Target', alpha=0.5)
plt.legend()

# Finally show all plotting
plt.show()
